# Copyright (c) 2020 cgsdfc
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT
"""
Unparse an asdl module.
"""

from typing import Sequence
from astlib import asdl
from pathlib import Path
from io import StringIO

ASDL_COMMENT_PREFIX = '--'
ASDL_MODULE_ATTR_COMMENT = 'comment'


class UnparsingVisitor(asdl.VisitorBase):
    """
    This visitor unparses the asdl tree.
    """
    def __init__(self):
        super().__init__()
        self.ostream = StringIO()

    def Unparse(self, node: asdl.AST) -> str:
        "Unparse an AST node."
        self.visit(node)
        return self.ostream.getvalue()

    def WriteComment(self, mod: asdl.Module):
        comment = getattr(mod, ASDL_MODULE_ATTR_COMMENT, None)
        if comment is None:
            return

        for line in comment.splitlines(keepends=True):
            self.ostream.write(ASDL_COMMENT_PREFIX + ' ' + line)
        self.ostream.write('\n')

    def visitModule(self, mod: asdl.Module):
        self.WriteComment(mod)
        self.ostream.write('module {}\n'.format(mod.name))
        self.ostream.write('{\n')
        for dfn in mod.dfns:
            self.ostream.write('\t')
            self.visit(dfn)
            self.ostream.write('\n\n')
        self.ostream.write('}\n')

    def visitType(self, type: asdl.Type):
        self.ostream.write('{} = '.format(type.name))
        self.visit(type.value)

    def _visitAttributes(self, attrs):
        if attrs:
            self.ostream.write('\t\t attributes ')
            self._visitFields(attrs)

    def visitSum(self, sum_: asdl.Sum):
        for i, ctor in enumerate(sum_.types):
            if i > 0:
                self.ostream.write('\t\t | ')
            self.visit(ctor)
            self.ostream.write('\n')
        self._visitAttributes(sum_.attributes)

    def _visitFields(self, fields: Sequence[asdl.Field]):
        if not fields:
            return

        def visitField(field: asdl.Field):
            if field.opt:
                suffix = '?'
            elif field.seq:
                suffix = '*'
            else:
                suffix = ''
            return '{0.type}{1} {0.name}'.format(field, suffix)

        val = '({})'.format(', '.join(map(visitField, fields)))
        self.ostream.write(val)

    def visitProduct(self, prod: asdl.Product):
        assert prod.fields
        self._visitFields(prod.fields)
        self.ostream.write('\n')
        self._visitAttributes(prod.attributes)

    def visitConstructor(self, ctor: asdl.Constructor):
        self.ostream.write(ctor.name)
        self._visitFields(ctor.fields)


def UnparseAsdlModule(mod: asdl.Module, output: Path):
    vis = UnparsingVisitor()
    with output.open('w') as f:
        result = vis.Unparse(mod)
        f.write(result)
